# 图像分类
import os
import json
from app import image_util
from flask import Blueprint, request, jsonify
from ultralytics import YOLO

# 创建蓝图
classify_bp = Blueprint('classify_bp', __name__)

# 获取当前脚本的绝对路径
script_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(script_dir)

# 加载多个 YOLO 模型并存储在字典中
#   models = {
#       'model1': YOLO('./runs/classify/train3/weights/best.pt')，
#       'model2': YOLO('./runs/classify/train3/weights/best1.pt')
#   }

# classify_000 :  通用分类模型 （  bd： 记磅单、car： 车辆照片、rkd： 入库单（入库明细）、 zxzm： 注销证明）
models = {
    'classify_000': YOLO(os.path.join(parent_dir, 'static', 'model', 'classify', 'best.pt'))
}


#
@classify_bp.route('/inference_image', methods=['POST', 'GET'])
def inference_image():
    print('api   /classify/inference_image')
    img_path = ''
    model_code = ''
    # 判断请求类型 获取请求参数
    if request.method == 'GET':
        # 获取 GET 请求参数
        img_path = request.args.get('imgPath')
        model_code = request.args.get('modelCode')
    elif request.method == 'POST':
        # 获取表单数据
        if request.form:
            data = request.form
            img_path = data.get('imgPath')
            model_code = data.get('modelCode')
        # 获取 JSON 数据
        elif request.json:
            data = request.json
            img_path = data.pop('imgPath')
            model_code = data.pop('modelCode')
    else:
        return "Unsupported request method"

    print(img_path)
    print(model_code)
    # 获取图片
    image = image_util.read_image(img_path)
    return inference_image_from_image(image, model_code)


def inference_image_from_image(image, model_code):
    """
    使用指定的分类模型对图片进行推理
    :param image:  推理图片
    :param model_code:  模型code
    :return:  推理结果
    """
    if model_code not in models:
        return jsonify({'error': f'Model {model_code} not found'}), 404

    # 使用指定的模型进行预测
    model = models[model_code]
    res = model(image)
    return json.dumps(parse_classify_results(res), indent=4)


def parse_classify_results(results):
    """
    解析分类结果，返回置信度最高的分类及其置信率。
    :param results: 列表，每个元素是一个 (class_name, confidence) 的元组
    :return: 字典，包含分类名称和置信率
    """
    if isinstance(results, list):
        # 如果 results 是一个列表，处理第一个元素
        results = results[0]
        # print(results)

    if not results:
        return {"error": "No classification results provided"}
    class_idx = results.probs.top1
    # 获取对应的类别名称和置信度
    class_name = results.names[class_idx]
    confidence = results.probs.top1conf.item()
    print(f"Class name: {class_name}, confidence: {confidence}")
    return {
        "classify": class_name,
        "confidence": confidence
    }
